package edu.harvard.s3xyback.edu.harvard.s3xyback.examples.crypto;

import java.io.File;
import java.io.FileInputStream;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.Signature;
import java.security.spec.PKCS8EncodedKeySpec;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

/**
 * 
 * @author jkennedy
 *
 */
public class CookieEncoderAESwRSA {

	// key sizes to generate
	private static final int RSA_KEYSIZE = 2048; //bits
	private static final int AES_KEYSIZE = 256; // bits
	private static final int SIG_SIZE = RSA_KEYSIZE/8; // bytes
	
	// algorithms to use
	private static final String RSA_STRING = "RSA/ECB/PKCS1Padding";
	private static final String AES_STRING = "AES/CBC/PKCS5Padding";
	private static final String SIG_STRING = "SHA1withRSA";	
	
	// Keys
	private PrivateKey rsaPrivateKey; // Private key used by AMS server
	private SecretKey aesSessionKey; // Secret key generated for this cookie
	
	// ciphers
	private Signature rsaSig;
	private Cipher rsaCipher;	
	private Cipher aesCipher;
	private KeyGenerator aesKeyGen;	
	
	
	/**
	 * Instantiate the decoder for use with a given private key. This key
	 * will be used for both encryption and signing.
	 * 
	 * @param key rsa Private key bytes
	 * @throws CryptoException
	 */
	public CookieEncoderAESwRSA(byte[] key) throws CryptoException {
		initKey(key);
	}
	
	
	/**
	 * Instantiate the encoder for use with a given private key. This key
	 * will be used for both encryption and signing. 
	 * 
	 * @param keyFile file that contains the Private key bytes
	 * @throws CryptoException
	 */
	public CookieEncoderAESwRSA(File keyFile) throws CryptoException { 
		try {
			
			byte[] key = new byte[(int)keyFile.length()];
			new FileInputStream(keyFile).read(key);
			initKey(key);
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
	}
	
	
	/**
	 * Parse the RSA key bytes and init the ciphers being used.
	 * 
	 * @param key
	 * @throws CryptoException
	 */
	private void initKey(byte[] key) throws CryptoException {
		try {
			
			// load Private key
			PKCS8EncodedKeySpec PrivateKeySpec = new PKCS8EncodedKeySpec(key);
			KeyFactory kf = KeyFactory.getInstance("RSA");
			rsaPrivateKey = kf.generatePrivate(PrivateKeySpec);			
		
			
			// load and init rsa cipher
			rsaCipher = Cipher.getInstance(RSA_STRING);
			rsaCipher.init(Cipher.WRAP_MODE, rsaPrivateKey);
			
			// load and init rsa signature
			rsaSig = Signature.getInstance(SIG_STRING);
			rsaSig.initSign(rsaPrivateKey);
			
			// init aes key generator
			aesKeyGen = KeyGenerator.getInstance("AES");
			aesKeyGen.init(AES_KEYSIZE);
			
			// load aes algorithm but can't init until initSession()
			aesCipher = Cipher.getInstance(AES_STRING);		
			
		} catch(Exception e) {
			throw new CryptoException(e);
		}
	}	
	
	/**
	 * 
	 * @return
	 */
	public String initSession() throws CryptoException {
		try {
			
			aesSessionKey = aesKeyGen.generateKey();
			byte[] iv = subset(aesSessionKey.getEncoded(), 0, 16); // shortcut to generate a random IV
			aesCipher.init(Cipher.ENCRYPT_MODE, aesSessionKey, new IvParameterSpec(iv));
			byte[] encKey = rsaCipher.wrap(aesSessionKey);
			return Base64.encodeToString(encKey);
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
		
	}
	
	/**
	 * 
	 * @param msg
	 * @return
	 */
	public String encrypt(String msg) throws CryptoException {
		try {
					
			// sign the plaintext with RSA
			byte[] msgBytes = msg.getBytes("UTF8");
			rsaSig.update(msgBytes);
			byte[] sigBytes = rsaSig.sign();		
			
			// encrypt the plaintext and signature with AES				
			byte[] encMsg = aesCipher.doFinal(concatBytes(sigBytes, msgBytes));		
			
			// return the encrypted session key and
			// encrypted plaintext concatenated together		
			return Base64.encodeToString(encMsg);
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
	}
	
	
	/**
	 * Combine two arrays together
	 * @param a1
	 * @param a2
	 * @return
	 */
	private static byte[] concatBytes(byte[] a1, byte[] a2) {
		byte[] concat = new byte[a1.length + a2.length];
		
		int c = 0;
		for (int i = 0; i < a1.length; i++) {
			concat[c++] = a1[i];
		}
		for (int i = 0; i < a2.length; i++) {
			concat[c++] = a2[i];			
		}
		
		return concat;
	}	
	
	/**
	 * Return a subset of the byte array as a new byte array
	 * @param bytes
	 * @param start
	 * @param end
	 * @return
	 */
	private static byte[] subset(byte[] bytes, int start, int end) {
		int length = end-start;
		if (length < 1) throw new IllegalArgumentException("Can not end before start");
		byte[] buff = new byte[length];
		for (int i = 0; i < length; i++) {
			buff[i] = bytes[i+start];
		}
		return buff;
	}	
	
}
